import torch
import cv2
from collections import defaultdict
import numpy as np

# ---------- 配置 ----------
VIDEO_IN = r'D:\Projects\person_count\yolov5\yolov5-master\campus_gate.mp4'
VIDEO_OUT = r'D:\Projects\person_count\yolov5\yolov5-master\out.avi'
YOLO_WEIGHTS = r'D:/Projects/person_count/yolov5/yolov5-master/yolov5s.pt'
LINE_X = 400        # 左侧判定线位置
LINE_X2 = 1050      # 右侧判定线位置
MIN_CONF = 0.5

# ---------- 加载模型 ----------
model = torch.hub.load('ultralytics/yolov5', 'custom', path=YOLO_WEIGHTS)
model.eval()
model.conf = MIN_CONF
model.classes = [0]

# ---------- 视频处理 ----------
cap = cv2.VideoCapture(VIDEO_IN)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
out = cv2.VideoWriter(VIDEO_OUT, cv2.VideoWriter_fourcc(*'XVID'), fps, (w, h))

# ---------- 跟踪 & 计数 ----------
track_id_counter = 0
tracked_objects = {}  # track_id: {'box': [x1, y1, x2, y2], 'cx': 上一帧中心x}
in_count1 = 0
out_count1 = 0
in_count2 = 0
out_count2 = 0

def compute_iou(box1, box2):
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    inter_area = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2]-box1[0]) * (box1[3]-box1[1])
    area2 = (box2[2]-box2[0]) * (box2[3]-box2[1])
    union_area = area1 + area2 - inter_area
    return inter_area / union_area if union_area > 0 else 0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame)
    dets = []

    # 获取YOLO检测结果
    for det in results.xyxy[0]:
        if len(det) == 6:
            x1, y1, x2, y2, conf, cls = det.tolist()
            if int(cls) == 0:
                dets.append([int(x1), int(y1), int(x2), int(y2)])

    updated_ids = set()
    new_tracked = {}

    for box in dets:
        x1, y1, x2, y2 = box
        cx = (x1 + x2) // 2
        matched_id = None
        max_iou = 0

        for tid, data in tracked_objects.items():
            iou = compute_iou(box, data['box'])
            if iou > 0.3 and iou > max_iou and tid not in updated_ids:
                matched_id = tid
                max_iou = iou

        # 如果找到了匹配的目标ID
        if matched_id is not None:
            prev_cx = tracked_objects[matched_id]['cx']
            
            # 处理左侧判定线（in_count1, out_count1）
            if prev_cx < LINE_X <= cx:  # 从左到右穿过左侧判定线
                in_count1 += 1
            elif prev_cx > LINE_X >= cx:  # 从右到左穿过左侧判定线
                out_count1 += 1

            # 处理右侧判定线（in_count2, out_count2）
            if prev_cx > LINE_X2 >= cx:  # 从右到左穿过右侧判定线
                in_count2 += 1
            elif prev_cx < LINE_X2 <= cx:  # 从左到右穿过右侧判定线
                out_count2 += 1
            
            new_tracked[matched_id] = {'box': box, 'cx': cx}
            updated_ids.add(matched_id)
        else:
            track_id_counter += 1
            new_tracked[track_id_counter] = {'box': box, 'cx': cx}
            updated_ids.add(track_id_counter)

    # 如果某些目标丢失，进行重新跟踪
    # 这里可以增加一个丢失检测，比如暂时丢失的目标，在几帧内保持原ID
    tracked_objects = new_tracked

    # ---------- 可视化 ----------
    for tid, data in tracked_objects.items():
        x1, y1, x2, y2 = data['box']
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 绿框
        cv2.putText(frame, f'ID{tid}', (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    # 绘制判定线
    cv2.line(frame, (LINE_X, 0), (LINE_X, h), (0, 0, 255), 2)   # 左侧判定线
    cv2.line(frame, (LINE_X2, 0), (LINE_X2, h), (255, 0, 0), 2)  # 右侧判定线

    # 显示进出人数
    cv2.putText(frame, f'IN1: {in_count1}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.putText(frame, f'OUT1: {out_count1}', (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    cv2.putText(frame, f'IN2: {in_count2}', (10, 110), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
    cv2.putText(frame, f'OUT2: {out_count2}', (10, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    out.write(frame)
    cv2.imshow('frame', frame)
    if cv2.waitKey(1) == 27:
        break

cap.release()
out.release()
cv2.destroyAllWindows()

# 最终统计
total_in = in_count1 + in_count2
total_out = out_count1 + out_count2
print(f'最终统计 → 进入食堂: {total_in}，离开食堂: {total_out}')
